# IK和IKV逻辑完全不同
# TODO 把参数放入单独的配置文件
# IK是直接计算结果，只要求出结果了，就不再迭代，每次直接给出最终q即可
# IKV是每次给速度，但是直到实际收敛了，再给出0
import pinocchio as pin
import numpy as np
from os.path import dirname, join
from scipy.spatial.transform import Rotation
import math




# pinocchio加载模型
# 参数>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
motor_num = 6


modelpath_dof7 = join(dirname(__file__), 'Dof7Arm.xml')
pin_model_dof7 = pin.RobotWrapper.BuildFromMJCF(modelpath_dof7)
site_id_dof7 = pin_model_dof7.model.getFrameId("site_Hand")

modelpath_dof6 = join(dirname(__file__), 'Dof6Arm_new.xml')
pin_model_dof6 = pin.RobotWrapper.BuildFromMJCF(modelpath_dof6)
site_id_dof6 = pin_model_dof6.model.getFrameId("site_Hand")


damping: float = 1e-3
diag_dof6 = damping * np.eye(6)
diag_dof7 = damping * np.eye(7)
qa_max = 5.0


dt_ik = 0.1 # IK步长
IT_MAX = 50 # IK步数
convergence_threshold = 1e-3
kv = 0.1

admittance_kf = 1e-3
# admittance_kf = 3e-3
admittance_kd = 1e-3
max_bias = 0.3


control_period = .001
control_time = float(30)
rad_per_sec = 100000 # 最大速度，每秒转角
max_dq = min(int(control_period*rad_per_sec), 10000) # 最大转角差


# 滤波参数
alpha = 0.02
#<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<


def percent_clip(x = np.zeros(motor_num),max = 0):
    if np.max(np.abs(x))>max:
        return x/np.max(np.abs(x))*max
    else:
        return x

def pos_admittance(k,x=np.zeros(3)):
    pos_bias_raw = k*x*np.abs(x)
    pos_bias = np.clip(pos_bias_raw, -max_bias, max_bias)
    print(f"pos_bias is {pos_bias}")
    return pos_bias


def get_interpoint(q_real, intikq_clip, max_delta_intq_eachloop):
    if not hasattr(get_interpoint, 'last_point'):
        get_interpoint.last_point = q_real



    if np.max(np.abs(intikq_clip - q_real)) > max_delta_intq_eachloop:
        delta_intq_new = max_delta_intq_eachloop/np.max(np.abs(intikq_clip - q_real)) * (intikq_clip - q_real)
    elif np.max(np.abs(intikq_clip - get_interpoint.last_point)) > max_delta_intq_eachloop:
        delta_intq_new = max_delta_intq_eachloop/np.max(np.abs(intikq_clip - get_interpoint.last_point)) * (intikq_clip - get_interpoint.last_point)

    else:
        delta_intq_new = (intikq_clip - get_interpoint.last_point)

    ctrl_intq = np.rint(delta_intq_new + get_interpoint.last_point)
    get_interpoint.last_point = ctrl_intq

    return ctrl_intq

def get_interpoint_l(q_real, intikq_clip, max_delta_intq_eachloop):
    if not hasattr(get_interpoint_l, 'last_point'):
        get_interpoint_l.last_point = q_real

    if np.max(np.abs(intikq_clip - q_real)) > max_delta_intq_eachloop:
        delta_intq_new = max_delta_intq_eachloop/np.max(np.abs(intikq_clip - q_real)) * (intikq_clip - q_real)
    elif np.max(np.abs(intikq_clip - get_interpoint_l.last_point)) > max_delta_intq_eachloop:
        delta_intq_new = max_delta_intq_eachloop/np.max(np.abs(intikq_clip - get_interpoint_l.last_point)) * (intikq_clip - get_interpoint_l.last_point)

    else:
        delta_intq_new = (intikq_clip - get_interpoint_l.last_point)

    ctrl_intq = np.rint(delta_intq_new + get_interpoint_l.last_point)
    get_interpoint_l.last_point = ctrl_intq

    return ctrl_intq



def get_interpoint_r(q_real, intikq_clip, max_delta_intq_eachloop):
    if not hasattr(get_interpoint_r, 'last_point'):
        get_interpoint_r.last_point = q_real


    if np.max(np.abs(intikq_clip - q_real)) > max_delta_intq_eachloop:
        delta_intq_new = max_delta_intq_eachloop/np.max(np.abs(intikq_clip - q_real)) * (intikq_clip - q_real)
    elif np.max(np.abs(intikq_clip - get_interpoint_r.last_point)) > max_delta_intq_eachloop:
        delta_intq_new = max_delta_intq_eachloop/np.max(np.abs(intikq_clip - get_interpoint_r.last_point)) * (intikq_clip - get_interpoint_r.last_point)
    else:
        delta_intq_new = (intikq_clip - get_interpoint_r.last_point)

    ctrl_intq = np.rint(delta_intq_new + get_interpoint_r.last_point)
    get_interpoint_r.last_point = ctrl_intq

    return ctrl_intq


def trans_intq2q_degree(intq):
    q_degree = intq/10000
    return q_degree

def trans_q_degree2intq(q_degree):
    intq = np.rint(q_degree*10000)
    return intq

def trans_intq2q(intq):
    q = intq*math.pi/180/10000
    return q

def trans_q2intq(q):
    intq = np.rint(q*10000*180/math.pi)
    return intq


def trans_intv2v(intv):
    v_real_float = intv*math.pi/15000
    return v_real_float

def trans_v2intv(v):
    intv = np.rint(v*15000/math.pi)
    return intv





def ik6(q_real_float, mocap_pos, mocap_quat, site_force = np.zeros(6)):
    # print(f"mocap_pos: {mocap_pos}")
    # print(f"mocap_quat: {mocap_quat}")
    # 判断是否ik_q是否已经满足要求
    if not hasattr(ik6, 'ik_qstar'):  # 初始化前次控制量
        ik6.ik_qstar = q_real_float # 从当前姿态开始迭代
    if not hasattr(ik6, 'converged'):
        ik6.converged = False
    pin.framesForwardKinematics(pin_model_dof6.model, pin_model_dof6.data, ik6.ik_qstar)
    pin.updateFramePlacements(pin_model_dof6.model, pin_model_dof6.data)


    # 求取base frame下的受力，并且转化为des_pos
    site_frame = pin_model_dof6.data.oMf[pin_model_dof6.model.getFrameId("site_Hand")]
    site_rotation = site_frame.rotation
    base_force = site_rotation.dot(site_force[:3])
    base_torque = site_rotation.dot(site_force[3:])
    



    des_pos = mocap_pos + pos_admittance(admittance_kf,base_force)
    # 获取目标位置SE3矩阵T_t
    mocap_rot = Rotation.from_quat([mocap_quat[1], mocap_quat[2], mocap_quat[3], mocap_quat[0]]).as_matrix()
    T_t = pin.SE3(mocap_rot, des_pos)
    T_bt = site_frame.actInv(T_t)



    # 定义的误差向量
    e_q = pin.log(T_bt).vector
    if np.linalg.norm(e_q) < convergence_threshold:
        ik6.converged = True
        print("Convergence ACHIEVED!")
        # print(f"e_q: {e_q}")
    else:
        ik6.converged = False
        print("Convergence FAILED!")
        # print(f"e_q: {e_q}")


    if not ik6.converged :
        
        pin.framesForwardKinematics(pin_model_dof6.model, pin_model_dof6.data, q_real_float+1e-20) # 这里添加一个非常小的小量，防止奇异
        pin.updateFramePlacements(pin_model_dof6.model, pin_model_dof6.data)

        ik_q = ik6.ik_qstar # 用于计算IK的关节q

        i = 0

        while True:



            # 求当前site的位姿T_b和目标位姿的相对差值矩阵T_bt
            T_bt = pin_model_dof6.data.oMf[pin_model_dof6.model.getFrameId("site_Hand")].actInv(T_t)
            # 定义的误差向量
            e_q = pin.log(T_bt).vector
            # 让误差较大时减轻计算量
            
            if np.linalg.norm(e_q) < convergence_threshold:
                ik6.converged = True
                ik6.ik_qstar = ik_q
                print("after cal Convergence ACHIEVED!")
                # print(f"after cal e_q: {e_q}")
                break
            if i>IT_MAX:
                ik6.converged = False
                ik6.ik_qstar = ik_q
                # print(f"after cal e_q: {e_q}")
                break
            
            i += 1

            # 求site_Hand的雅克比J_b
            J_b = pin.computeFrameJacobian(
                pin_model_dof6.model, 
                pin_model_dof6.data, 
                ik_q, 
                pin_model_dof6.model.getFrameId("site_Hand"), 
                pin.ReferenceFrame.LOCAL_WORLD_ALIGNED  # 根据需求选择坐标系
            )
            
            J_l = -pin.Jlog6(T_bt.inverse())
            # 接下来会这里会利用J_error v^{*} = -e(q)来求v^{*}
            J_error = J_l.dot(J_b) # 误差Jacobian

            # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
            v_star = np.linalg.solve(J_error.T.dot(J_error) + diag_dof6, -J_error.T.dot(e_q))
            # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

            ik_q = pin.integrate(pin_model_dof6.model, ik_q, v_star*dt_ik)
            pin.forwardKinematics(pin_model_dof6.model, pin_model_dof6.data, ik_q)
            pin.updateFramePlacements(pin_model_dof6.model, pin_model_dof6.data)

    # print(f"ik_qstar: {ik6.ik_qstar}")

    return ik6.ik_qstar



def ik7(q_real_float, mocap_pos, mocap_quat, qv_real_float = np.zeros(7), site_force = np.zeros(6)):
    # print(f"mocap_pos: {mocap_pos}")
    # print(f"mocap_quat: {mocap_quat}")
    # 判断是否ik_q是否已经满足要求
    if not hasattr(ik7, 'ik_qstar'):  # 初始化前次控制量
        ik7.ik_qstar = q_real_float # 从当前姿态开始迭代
    if not hasattr(ik7, 'converged'):
        ik7.converged = False
    pin.framesForwardKinematics(pin_model_dof7.model, pin_model_dof7.data, ik7.ik_qstar)
    pin.forwardKinematics(pin_model_dof7.model, pin_model_dof7.data, ik7.ik_qstar, qv_real_float)
    pin.updateFramePlacements(pin_model_dof7.model, pin_model_dof7.data)


    # 求取base frame下的受力，并且转化为des_pos
    site_frame = pin_model_dof7.data.oMf[site_id_dof7]
    site_rotation = site_frame.rotation
    base_force = site_rotation.dot(site_force[:3])
    base_torque = site_rotation.dot(site_force[3:])
    print(f"site force is {site_force}")

    velocity_base = pin.getFrameVelocity(pin_model_dof7.model, pin_model_dof7.data, site_id_dof7)
    print(f"velocity_base is {velocity_base}")

    # 引入qv_real
    des_pos = mocap_pos + pos_admittance(admittance_kf,base_force - admittance_kd*velocity_base.linear)
    # 获取目标位置SE3矩阵T_t
    mocap_rot = Rotation.from_quat([mocap_quat[1], mocap_quat[2], mocap_quat[3], mocap_quat[0]]).as_matrix()
    T_t = pin.SE3(mocap_rot, des_pos)
    T_bt = site_frame.actInv(T_t)

    # 定义的误差向量
    e_q = pin.log(T_bt).vector
    if np.linalg.norm(e_q) < convergence_threshold:
        ik7.converged = True
        print("Convergence ACHIEVED!")
        # print(f"e_q: {e_q}")
    else:
        ik7.converged = False
        print("Convergence FAILED!")
        # print(f"e_q: {e_q}")


    if not ik7.converged :
        
        pin.framesForwardKinematics(pin_model_dof7.model, pin_model_dof7.data, q_real_float+1e-20) # 这里添加一个非常小的小量，防止奇异
        pin.updateFramePlacements(pin_model_dof7.model, pin_model_dof7.data)

        ik_q = ik7.ik_qstar # 用于计算IK的关节q

        i = 0

        while True:

            pin_model_dof7.data.oMf[site_id_dof7]

            # 求当前site的位姿T_b和目标位姿的相对差值矩阵T_bt
            T_bt = pin_model_dof7.data.oMf[site_id_dof7].actInv(T_t)
            # 定义的误差向量
            e_q = pin.log(T_bt).vector
            # 让误差较大时减轻计算量
            
            if np.linalg.norm(e_q) < convergence_threshold:
                ik7.converged = True
                ik7.ik_qstar = ik_q
                print("after cal Convergence ACHIEVED!")
                # print(f"after cal e_q: {e_q}")
                break
            if i>IT_MAX:
                ik7.converged = False
                ik7.ik_qstar = ik_q
                # print(f"after cal e_q: {e_q}")
                break

            i += 1

            # 求site_Hand的雅克比J_b
            J_b = pin.computeFrameJacobian(
                pin_model_dof7.model, 
                pin_model_dof7.data, 
                ik_q, 
                pin_model_dof7.model.getFrameId("site_Hand"), 
                pin.ReferenceFrame.LOCAL_WORLD_ALIGNED  # 根据需求选择坐标系
            )
            
            J_l = -pin.Jlog6(T_bt.inverse())
            # 接下来会这里会利用J_error v^{*} = -e(q)来求v^{*}
            J_error = J_l.dot(J_b) # 误差Jacobian

            # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
            v_star = np.linalg.solve(J_error.T.dot(J_error) + diag_dof7, -J_error.T.dot(e_q))
            # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

            ik_q = pin.integrate(pin_model_dof7.model, ik_q, v_star*dt_ik)
            pin.forwardKinematics(pin_model_dof7.model, pin_model_dof7.data, ik_q)
            pin.updateFramePlacements(pin_model_dof7.model, pin_model_dof7.data)

    # print(f"ik_qstar: {ik.ik_qstar}")

    return ik7.ik_qstar



# def ikv(q_real_float, mocap_pos, mocap_quat):

#     # 判断是否ik_q是否已经满足要求
#     if not hasattr(ikv, 'ikv_qstar'):  # 初始化前次控制量
#         ikv.ikv_qstar = q_real_float
#     if not hasattr(ikv, 'converged'):
#         ikv.converged = False

#     pin.framesForwardKinematics(pin_model.model, pin_model.data, ikv.ikv_qstar+1e-20) # 这里添加一个非常小的小量，防止奇异
#     pin.updateFramePlacements(pin_model.model, pin_model.data)

#     # 获取目标位置SE3矩阵T_t
#     mocap_rot = Rotation.from_quat([mocap_quat[1], mocap_quat[2], mocap_quat[3], mocap_quat[0]]).as_matrix()
#     T_t = pin.SE3(mocap_rot, mocap_pos)
#     T_bt = pin_model.data.oMf[pin_model.model.getFrameId("site_Hand")].actInv(T_t)
#     # 定义的误差向量
#     e_q = pin.log(T_bt).vector

#     if np.linalg.norm(e_q) < convergence_threshold:
#         ikv.converged = True
#         print("Convergence ACHIEVED!")
#         # print(f"e_q: {e_q}")
#     else:
#         ikv.converged = False
#         print("Convergence FAILED!")
#         # print(f"e_q: {e_q}")

#     if not ikv.converged :
        
#         pin.framesForwardKinematics(pin_model.model, pin_model.data, q_real_float+1e-20) # 这里添加一个非常小的小量，防止奇异
#         pin.updateFramePlacements(pin_model.model, pin_model.data)

#         i = 0

#         while True:

#             # 求当前site的位姿T_b和目标位姿的相对差值矩阵T_bt
#             T_bt = pin_model.data.oMf[pin_model.model.getFrameId("site_Hand")].actInv(T_t)
#             # 定义的误差向量
#             e_q = pin.log(T_bt).vector

#             # 求site_Hand的雅克比J_b
#             J_b = pin.computeFrameJacobian(
#                 pin_model.model, 
#                 pin_model.data, 
#                 ikv.ikv_qstar, 
#                 pin_model.model.getFrameId("site_Hand"), 
#                 pin.ReferenceFrame.LOCAL_WORLD_ALIGNED  # 根据需求选择坐标系
#             )
            
#             J_l = -pin.Jlog6(T_bt.inverse())
#             # 接下来会这里会利用J_error v^{*} = -e(q)来求v^{*}
#             J_error = J_l.dot(J_b) # 误差Jacobian

#             # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
#             v_star = np.linalg.solve(J_error.T.dot(J_error) + diag, -J_error.T.dot(e_q))
#             # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

#             ikv.ikv_qstar = pin.integrate(pin_model.model, ikv.ikv_qstar, v_star*dt_ik)
#             pin.forwardKinematics(pin_model.model, pin_model.data, ikv.ikv_qstar)
#             pin.updateFramePlacements(pin_model.model, pin_model.data)

#             i += 1

#             if np.linalg.norm(e_q) < convergence_threshold:
#                 ikv.converged = True
#                 # print("after cal Convergence ACHIEVED!")
#                 # print(f"after cal e_q: {e_q}")
#                 break
#             if i>IT_MAX:
#                 ikv.converged = False
#                 # print(f"after cal e_q: {e_q}")
#                 break

#     ikv_star = kv*(ikv.ikv_qstar - q_real_float)

#     return ikv_star

if __name__ == "__main__":
    print("Hello World!")
    q_real_float = np.array([0,0,0,0,0,0])+damping
    mocap_pos = np.array([-0.1,0,1])
    mocap_quat = np.array([1,0,0,0])
    ikq_star = ik6(q_real_float, mocap_pos, mocap_quat)
    print(f"type of ikq_star: {type(ikq_star)}")
    q_list = ikq_star.tolist()
    print(f"tyepe of q_list: {type(q_list)}")
    print(f"q_list: {q_list}")